#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

# !/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
from tqdm import tqdm

# valid_chrom = set(
#     ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15", "16", "17", "18", "19", "20",
#      "21", "22"])

client = MongoClient("mongodb://han:han2019@mongo-han.handb:27017")
collection1 = client["hanvcf"]["all_20220612"]
collection2 = client["hanvcf"]["snvs_20220612"]

infile = open("variation_info_GRCh37_20220705.txt", "r")
outfile = open("variation_info_GRCh37_addhan_20220705.tsv", "w")
outfile.write(infile.readline().rstrip("\n") + "\tPggHan Sample Size\tPggHan Ref Allele\tPggHan Alt Allele\n")
outfile.flush()
for line in tqdm(infile):
    line = line.rstrip("\n")
    splitline = line.split("\t")
    rsid1 = splitline[0]
    rsid2 = splitline[1]
    rsid_list = rsid1.split("|")
    if rsid2 != "NA":
        rsid_list.append(rsid2)
    rsid_list = list(set(rsid_list))
    result2 = collection2.find_one({"rsid": {"$in": rsid_list}}, {"chrom": 1, "pos": 1})
    if result2 is not None:
        chrom = result2["chrom"]
        pos = result2["pos"]
        result_count = collection1.count_documents({"chrom": chrom, "pos": pos})
        if result_count == 0:
            outfile.write(line + "\tnull\tnull\tnull\n")
        elif result_count == 1:
            result = collection1.find_one({"chrom": chrom, "pos": pos})
            ref = result["ref"]
            alt = result["alt"]
            individual_count = 0
            ref_count = 0
            alt_count = 0
            for item in result["region"]:
                # print("%s\t%s\t%d" % (chrom, pos, item["individual_count"]))
                if item["individual_count"] is not None:
                    individual_count += item["individual_count"]
                    ref_count += item["rr"] * 2
                    ref_count += item["ra"]
                    alt_count += item["ra"]
                    alt_count += item["aa"] * 2
            if individual_count != 0:
                individual_count = individual_count * 2
                alt_freq = float(alt_count) / float(individual_count)
                ref_freq = float(ref_count) / float(individual_count)
                ref_str = "%s=%f" % (ref, ref_freq)
                alt_str = "%s=%f" % (alt, alt_freq)

                outfile.write(line + "\t%d\t%s\t%s\n" % (individual_count, ref_str, alt_str))
            else:
                outfile.write(line + "\tnull\tnull\tnull\n")
        else:
            results = collection1.find({"chrom": chrom, "pos": pos})
            individual_count = 0
            ref_count = 0
            ref = ""
            alt_list = {}
            for result in results:
                ref = result["ref"]
                alt = result["alt"]
                alt_count = 0
                for item in result["region"]:
                    if item["individual_count"] is not None:
                        individual_count += item["individual_count"]
                        ref_count += item["rr"] * 2
                        ref_count += item["ra"]
                        alt_count += item["ra"]
                        alt_count += item["aa"] * 2
                alt_list[alt] = alt_count
            if individual_count != 0:
                individual_count = individual_count * 2
                ref_freq = float(ref_count) / float(individual_count)
                alt_str_list = []
                for alt in alt_list:
                    alt_freq = float(alt_list[alt]) / float(individual_count)
                    alt_str_list.append("%s=%f" % (alt, alt_freq))
                ref_str = "%s=%f" % (ref, ref_freq)
                alt_str = "|".join(alt_str_list)
                outfile.write(line + "\t%d\t%s\t%s\n" % (individual_count, ref_str, alt_str))
            else:
                outfile.write(line + "\tnull\tnull\tnull\n")
    else:
        outfile.write(line + "\tnull\tnull\tnull\n")

    outfile.flush()
infile.close()
outfile.close()
